package org.linloong.config.ai;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class DifyModelManager {

    private final Map<String, WebClient> webClientMap = new ConcurrentHashMap<>();
    private final Map<String, WebClient> fileWebClientMap = new ConcurrentHashMap<>();
    private final MultiDifyConfig multiDifyConfig;

    public DifyModelManager(MultiDifyConfig multiDifyConfig) {
        this.multiDifyConfig = multiDifyConfig;
        initializeWebClients();
        initializeFileWebClients();
    }

    private void initializeWebClients() {
        if (multiDifyConfig.getModels() != null) {
            multiDifyConfig.getModels().forEach((modelName, config) -> {
                WebClient webClient = WebClient.builder()
                        .baseUrl(config.getBaseUrl())
                        .codecs(configurer -> configurer.defaultCodecs().maxInMemorySize(100 * 1024 * 1024))
                        .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                        .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + config.getApiKey())
                        .build();

                webClientMap.put(modelName, webClient);
                log.info("初始化 Dify 模型: {} - {}", modelName, config.getDescription());
            });
        }
    }
    private void initializeFileWebClients() {
        if (multiDifyConfig.getModels() != null) {
            multiDifyConfig.getModels().forEach((modelName, config) -> {
                WebClient webClient = WebClient.builder()
                        .baseUrl(config.getBaseUrl())
                        .codecs(configurer -> configurer.defaultCodecs().maxInMemorySize(10 * 1024 * 1024))
                        .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + config.getApiKey())
                        .build();

                fileWebClientMap.put(modelName, webClient);
                log.info("初始化 Dify 模型: {} - {}", modelName, config.getDescription());
            });
        }
    }
    public WebClient getWebClient(String modelName) {
        WebClient webClient = webClientMap.get(modelName);
        if (webClient == null) {
            throw new IllegalArgumentException("不支持的模型: " + modelName);
        }
        return webClient;
    }

    // 在配置类中添加

    public WebClient getFileWebClient(String modelName) {
        WebClient webClient = fileWebClientMap.get(modelName);
        if (webClient == null) {
            throw new IllegalArgumentException("不支持的模型: " + modelName);
        }
        return webClient;
    }



    public MultiDifyConfig.ModelConfig getModelConfig(String modelName) {
        MultiDifyConfig.ModelConfig config = multiDifyConfig.getModels().get(modelName);
        if (config == null) {
            throw new IllegalArgumentException("不支持的模型: " + modelName);
        }
        return config;
    }

    public Map<String, String> getAvailableModels() {
        Map<String, String> modelInfo = new ConcurrentHashMap<>();
        multiDifyConfig.getModels().forEach((name, config) -> {
            modelInfo.put(name, config.getDescription());
        });
        return modelInfo;
    }
}